package com.leeue.aop;

import com.google.common.util.concurrent.RateLimiter;
import com.leeue.annotation.ExtRateLimiter;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * @author liyue
 * @date 2019-08-12 16:34
 */
@Aspect
@Component
public class RateLimiterAop {

    //每一个请求 地址 只有一个 RateLimiter
    private Map<String, RateLimiter> rateLimiterMap = new HashMap<>();

    //1.定义切入点
    @Pointcut("execution( public  * com.leeue.api.*.*(..))")
    public void rlAop() {

    }
    //2.使用前置通知，看方法上是否有 ExtRateLimiter 注解，如果有就要进行限流操作

    @Around("rlAop()")
    public Object doBefore(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {

        //1.获取当前拦截的方法
        Method method = this.getSignature(proceedingJoinPoint);
        //2.判断是否有ExRateLimiter注解
        ExtRateLimiter extRateLimiter = method.getDeclaredAnnotation(ExtRateLimiter.class);
        if (extRateLimiter == null) {
            //直接报错
            return null;
        }
        //3.获取当前请求的url地址
        String uri = this.getRequestURI();
        if (StringUtils.isEmpty(uri)) {
            return null;
        }
        //判断当前Map中是否存在RateLimiter
        RateLimiter rateLimiter = null;
        if (rateLimiterMap.containsKey(uri)) {
            rateLimiter = rateLimiterMap.get(uri);
        } else {
            rateLimiter = RateLimiter.create(extRateLimiter.permitsPerSecond());
            rateLimiterMap.put(uri, rateLimiter);
        }

        //获取桶中令牌
        boolean tryAcquire = rateLimiter.tryAcquire(extRateLimiter.timeout(), TimeUnit.MILLISECONDS);
        if (!tryAcquire) {
            //如果不能获取到令牌就走服务降级方法
            this.fallback();
        }
        //否则就放行
        return proceedingJoinPoint.proceed();
    }


    /**
     * 服务降级方法
     */
    private void fallback() {

        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletResponse response = attributes.getResponse();
        response.setHeader("Content-type", "text/html;charset=UTF-8");
        PrintWriter writer = null;
        try {
            writer = response.getWriter();
            writer.println("服务繁忙,请稍后重试");
            System.out.println("走服务降级了--------");
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            writer.close();
        }
        // 在AOP 编程中获取响应
    }

    /**
     * 获取到URI地址
     *
     * @return
     */
    private String getRequestURI() {
        return getRequest().getRequestURI();
    }

    private HttpServletRequest getRequest() {
        //获取到请求的url
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        return attributes.getRequest();
    }

    /**
     * 获取到当前拦截的方法
     *
     * @param proceedingJoinPoint
     * @return
     */
    private Method getSignature(ProceedingJoinPoint proceedingJoinPoint) {
        MethodSignature signature = (MethodSignature) proceedingJoinPoint.getSignature();
        return signature.getMethod();
    }


}
